'''
构建pipeline 的gradio

img1【输入】             depth_img【输出】                restored_depth_img【输出】
btn_get_depth 【读取img1作为输入输出到depth_img】       btn_restored_depth 【读取depth_img 输出到 restored_depth_img】
img2【输入】 output_img【输出】
btn_get_res【读取img2和restored_depth_img输出到 output_img】

使用  
python xxxx.py  cuda_id  port
比如
python gradio_pipeline_for_depth_control.py 4 20021
'''

import argparse

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='示例脚本：接收CUDA设备和端口参数')
    
    # 添加参数
    parser.add_argument(
        '-c',
        '--cuda_id', 
        type=str,
        required=True,  # 必须传入
        help='CUDA设备ID，例如 "0" 或 "0,1"（字符串类型）'
    )
    parser.add_argument(
        '-p',
        '--port',
        type=int,
        default=20025,  # 默认值
        help='端口号（整数类型，默认8000）'
    )
    
    # 解析参数
    args = parser.parse_args()
    return args

args = parse_args()

import os,sys
os.environ['CUDA_VISIBLE_DEVICES']=args.cuda_id
import util_for_huggingface

import gradio as gr
import numpy as np
from PIL import Image,ImageOps
import cv2,requests,io,base64,torch

from diffusers import FluxControlPipeline, FluxPriorReduxPipeline
from util_flux import process_img_1024,vertical_concat_images,horizontal_concat_images

from image_gen_aux import DepthPreprocessor
# from util_mask import load_yolo,get_mask_by_yolo
from demo_rmbg import load_rmbg,get_mask_by_rmbg

from MODEL_CKP import FLUX_DEPTH,FLUX_REDUX,DEPTH_PREDCITION
depth_processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                    FLUX_REDUX, 
                                    torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxControlPipeline.from_pretrained(FLUX_DEPTH, torch_dtype=torch.bfloat16).to("cuda")
model = load_rmbg()

cur_in1 = None
cur_in2 = []
out = []

ori_depth = None
now_depth = None

save_dir = '/mnt/nas/shengjie/datasets/dataset_depth_control'
save_depth_dir = '/mnt/nas/shengjie/datasets/dataset_depth_control_depth'
os.makedirs( save_dir,exist_ok=True )
os.makedirs( save_depth_dir,exist_ok=True )

# 模拟处理函数 - 实际使用时替换为你的真实处理逻辑
# get depth 之后 再用mask 限制一下 再输出
def get_depth(input_img):
    """模拟深度图生成"""
    if input_img is None:
        return None
    # 这里用灰度图模拟深度图
    # depth = cv2.cvtColor(input_img, cv2.COLOR_RGB2GRAY)

    global cur_in1

    input_img = process_img_1024('',img_pil=input_img)

    cur_in1 = input_img

    # extra 处理
    # 输入 input_img => mask
    # mask + depth => output
    mask_pil,no_bg_img = get_mask_by_rmbg(model , input_img)
    
    img_pil = input_img
    
    depth_pil = depth_processor(img_pil)[0].convert('RGB')

    if mask_pil is not None:
        mask_pil = process_img_1024(mask_pil,pad_color=(0,0,0)).convert('L')
        # mask_np = np.array(mask_pil)

        black_mask = Image.new( 'RGB' , depth_pil.size ,0 )
        masked_depth = Image.composite( black_mask ,depth_pil , ImageOps.invert(mask_pil) ) # 正常是 黑留白遮 ops 白留黑遮
    
        # mask_pil depth_pil => 高斯模糊
        depth_pil_blur = replace_lama_cleaner( mask_pil , depth_pil )
        masked_depth_blur = Image.composite( black_mask ,depth_pil_blur , ImageOps.invert(mask_pil) ) # 正常是 黑留白遮 ops 白留黑遮
    
    else:
        masked_depth = depth_pil

    

    return masked_depth_blur.convert('RGB')

def image_to_bytes(image):
    """将PIL Image转换为字节流"""
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='JPEG')
    return img_byte_arr.getvalue()

from PIL import ImageFilter
def replace_lama_cleaner(mask_image,depth_pil):
    # get inverse mask_image


    # 1. 提取边缘信息，生成mask
    # 转为灰度
    gray = mask_image.convert("L")
    # 使用FIND_EDGES滤波器
    edges = gray.filter(ImageFilter.FIND_EDGES)
    # 增强边缘对比度
    edges = ImageOps.autocontrast(edges)
    # 二值化，得到mask
    edge_np = np.array(edges)
    threshold = 30  # 可调整
    mask_np = (edge_np > threshold).astype(np.uint8) * 255
    mask = Image.fromarray(mask_np, mode="L")
    # tmp_mask = mask

    # 对mask膨胀一点，让边缘更宽
    mask = mask.filter(ImageFilter.MaxFilter(7))

    # 2. 对边缘不模糊，对边缘内部做模糊处理
    # 先对整张图做高斯模糊
    blurred = depth_pil.filter(ImageFilter.GaussianBlur(radius=20))

    # mask反转，边缘为0，内部为255
    inv_mask = ImageOps.invert(mask)

    # pdb.set_trace()

    # 合成：边缘用原图，内部用模糊图
    #  mask_image_pil 的 白色部分 填充为 depth_pil 中 对应区域的 颜色
    mask_np = np.array(gray) > 200 # 在第三维度复制自己 => （h,w,ch)
    depth_np = np.array( depth_pil )
    mask_3d = np.repeat(mask_np[:, :, np.newaxis], 3, axis=2)
    mask_image_merged = np.zeros_like(depth_np) 
    mask_image_merged[mask_3d] = depth_np[mask_3d]
    mask_image_merged = Image.fromarray( mask_image_merged )
    result = Image.composite(blurred, mask_image_merged, inv_mask)

    # 保存
    horizontal_concat_images([
        mask_image , mask, inv_mask, 
        mask_image_merged , result
    ]).save("tmp_blur.jpg")
    return result


def process_image(input_image, mask_image):
    """
    处理图片并返回修复后的结果
    
    参数:
        input_image: PIL.Image对象，原始图片
        mask_image: PIL.Image对象，mask图片（白色表示要修复的区域）
        返回: 修复后的PIL.Image对象
    """
    # input_image = process_img_1024('',img_pil=input_image)
    # mask_image = process_img_1024('',img_pil=mask_image)

    # 准备要发送的数据
    files = {
        'image': ('input.jpg', image_to_bytes(input_image), 'image/jpeg'), # depth 
        'mask': ('mask.jpg', image_to_bytes(mask_image), 'image/jpeg') # 手动涂抹的mask
    }
    # 发送POST请求
    api_url = 'http://localhost:20021/api/inpaint'
    response = requests.post(api_url, files=files)
    if response.status_code == 200:
        # 解析返回的base64图像
        result_data = response.json()
        img_data = base64.b64decode(result_data['result'])
        return Image.open(io.BytesIO(img_data))
    else:
        raise Exception(f"API请求失败: {response.text}")

def restore_depth(depth_img):
    """模拟深度图修复"""
    if depth_img is None:
        return None
    input_image = depth_img["image"]
    mask_image = depth_img["mask"]

    input_image = process_img_1024('',img_pil=depth_img['image'])
    mask_image = process_img_1024('',img_pil=depth_img['mask']).convert("L") 
    restored_depth =  process_image( input_image , mask_image )
    
    global ori_depth,now_depth
    ori_depth = input_image
    now_depth = restored_depth
    
    return restored_depth

def get_result(img2, restored_depth, alpha = 0.3 ,control_alpha = 0.2):
    """模拟最终图像合成"""
    if img2 is None or restored_depth is None:
        return None
    # 这里用简单的混合模拟合成
    # restored_depth = cv2.cvtColor(restored_depth, cv2.COLOR_GRAY2RGB)
    # result = cv2.addWeighted(img2, 0.7, restored_depth, 0.3, 0)
    redux_image =  process_img_1024( '' , img_pil=img2 )
    control_image =  process_img_1024( '' , img_pil=restored_depth )

    

    with torch.no_grad():
        # 这里试试不用衣服的redux相加
        
        # tmp_path = 'tmp_other.jpg'
        # prompt_emb,pooled_prompt_emb = pipe_prior_redux(control_image,
        #                                                 return_dict=False) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 
        prompt_emb,pooled_prompt_emb = pipe_prior_redux(redux_image,return_dict=False) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 
        prompt_emb[:,-25:,:]=0
        # prompt_emb2[:,0::2,:]=0

        # prompt_emb = (1-alpha) * prompt_emb + alpha * prompt_emb2
        prompt_emb = alpha * prompt_emb

        image = pipe(
            # prompt=prompt,
            control_image=control_image,
            control_image2=control_image,
            control_alpha=control_alpha,
            height=control_image.height,
            width=control_image.width,
            num_inference_steps=20,
            guidance_scale=4.5,
            # generator=torch.Generator().manual_seed(42),
            # **pipe_prior_output,
            prompt_embeds=prompt_emb,
            pooled_prompt_embeds=pooled_prompt_emb,
        ).images[0]

        mask_pil , no_bg_img =   get_mask_by_rmbg( model , image )
        white_mask = Image.new( 'RGB' , image.size , 'white'  )
        masked_img = Image.composite( white_mask , image , ImageOps.invert(mask_pil) ) # 正常是 黑留白遮 ops 白留黑遮
    
    return masked_img

with gr.Blocks(title="图像处理Pipeline") as demo:
    gr.Markdown("# 图像处理Pipeline")
    gr.Markdown("上传图片，按步骤处理")
    
    height = 512

    with gr.Row():
        with gr.Column():
            # 第一部分：深度图生成与修复
            img1 = gr.Image(label="输入图像1", 
                            type="pil",height=height)
            btn_get_depth = gr.Button("生成深度图")

            depth_img = gr.Image(label="深度图", 
                                 type="pil",
                                 tool="sketch",
                                 height=height)
            
            btn_restored_depth = gr.Button("修复深度图")

            restored_depth_img = gr.Image(label="修复后的深度图", 
                                          type="pil", 
                                          interactive=False,
                                          height=height)
        
        with gr.Column():
            # 第二部分：最终图像合成
            img2 = gr.Image(label="输入图像2", type="pil",
                            height=height)
            # 新增：多文件上传组件
            file_upload = gr.Files(label="批量上传图片", file_types=["image"])

            # 设置两个可以滑动的 slide，区间在 [0,1]
            # alpha control_alpha 
            # 添加两个滑动条
            with gr.Row():
                alpha = gr.Slider(0, 1, value=0.5, label="风格比重")
                control_alpha = gr.Slider(0, 1, value=0.5, label="控制depth比重")
            

            btn_get_res = gr.Button("生成最终结果")
            btn_batch_process = gr.Button("批量处理所有上传图片")  # 新增批量处理按钮
            
            output_img = gr.Image(label="输出图像", type="pil", 
                                  interactive=False,
                                  height=height)

            batch_output_gallery = gr.Gallery(label="批量处理结果",
                                              rows=4,)  # 新增画廊显示批量结果
            
            
            ## todo list
            ## 新增 上传多个图片文件
            ## btn 处理所有的图片文件 使用 get_result 函数，循环调用传入图片文件，仅替换img2的输入即可

            def get_test_imgs():
                dir = '/mnt/nas/shengjie/datasets/images_sampled_1k'
                with open( os.path.join(dir , 'names.txt') ) as f:
                    names = f.readlines()
                return [ 
                    os.path.join( dir , name.strip() ) 
                    for name in names
                    if name.strip().endswith( '.jpg' )
                ]

            # 新增批量处理函数
            def batch_process(depth_img, alpha, control_alpha):
                img_paths = get_test_imgs()
                # files = [i_p for i_p in img_paths ]

                global cur_in1
                global out
                
                global ori_depth
                global now_depth
                # cur_in2 = []
                out = []

                results = []
                import time
                current_millis = int(time.time() * 1000)
                
                ## 另外save [ori_depth , now_depth] 
                depth_o = horizontal_concat_images( [
                    process_img_1024('',img_pil=ori_depth , pad_color=(0,0,0)), 
                    process_img_1024('',img_pil=now_depth , pad_color=(0,0,0))
                ] )
                # depth_o.save( os.path.join( save_depth_dir , f'{current_millis}.jpg' ) )
                
                i = 0
                for file in img_paths[:2]:
                    # 读取上传的图片文件
                    img = Image.open(file)
                    
                    # cur_in2.append( process_img_1024( '' , img_pil=img ) )

                    # 调用相同的处理函数
                    result = get_result(img, depth_img, alpha, control_alpha)
                    results.append(result)

                    try:
                        o = horizontal_concat_images( [ 
                            cur_in1 , process_img_1024( '' , img_pil=img ) , result,
                        ] ,  gap=0 , line_width=0, )
                        out.append( o )
                        # o.save(  os.path.join(save_dir, f'{current_millis}_{i}.jpg') )


                        i+=1
                    except :
                        pdb.set_trace()



                # 获取当前时间戳（秒.毫秒）
                return results[:10]

            # 批量处理按钮点击事件
            btn_batch_process.click(
                fn=batch_process,
                inputs=[restored_depth_img, alpha, control_alpha],
                outputs=batch_output_gallery
            )

            download_btn = gr.Button("打包下载所有图片")

            import zipfile,pdb
            from tempfile import mktemp
            def create_zip():
                global out
                images = out
                # 创建临时ZIP文件
                zip_path = mktemp('.zip')
                
                with zipfile.ZipFile(zip_path, 'w') as zipf:
                    for i, img in enumerate(images):
                        # 临时保存图片
                        img_path = mktemp('.jpg')
                        # pdb.set_trace()
                        # img = Image.open(img['name'])
                        img.save(img_path)
                        # 添加到ZIP
                        zipf.write(img_path, f"result_{i+1}.jpg")
                        os.remove(img_path)
                
                return zip_path

            download_btn.click(
                fn=create_zip,
                inputs=None,
                outputs=gr.File(label="下载ZIP包", type="file")
            )
            
    
    # 第一部分处理流程
    btn_get_depth.click(
        fn=get_depth,
        inputs=img1,
        outputs=depth_img
    )
    
    btn_restored_depth.click(
        fn=restore_depth,
        inputs=depth_img,
        outputs=restored_depth_img
    )
    
    # 第二部分处理流程
    btn_get_res.click(
        fn=get_result,
        inputs=[img2, restored_depth_img,alpha,control_alpha],
        outputs=output_img
    )
    
    # 示例数据
    # gr.Examples(
    #     examples=[
    #         ["examples/img1.jpg", "examples/img2.jpg"],
    #     ],
    #     inputs=[img1, img2],
    #     outputs=[depth_img, restored_depth_img, output_img],
    #     fn=lambda img1, img2: (
    #         get_depth(img1),
    #         restore_depth(get_depth(img1)),
    #         get_result(img2, restore_depth(get_depth(img1)))
    #     ),
    #     cache_examples=True,
    #     label="示例图片"
    # )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=args.port)